use std::sync::Arc;

use async_trait::async_trait;
use bytestring::ByteString;
use datafusion::arrow::array::RecordBatch;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::prelude::{DataFrame, SessionContext};
use datafusion::sql::parser::Statement;
use either::Either;
use sqlparser::ast::{Ident, ObjectName, ObjectType, ShowStatementFilter, Statement as SQLStatement};

use common::{Error, LResult};

pub type DBRef = Arc<dyn DB>;

pub type Data = Either<DataFrame, RecordBatch>;

#[async_trait]
pub trait DB: Send + Sync {
    async fn insert(&self, line: ByteString) -> LResult<()>;

    async fn query(&self, sql: &str) -> LResult<Data> {
        let ctx = self.session_context();
        let stmt = ctx.state().sql_to_statement(sql, "mysql").map_err(|e| Error::ExecutionErr(format!(
            "parse sql: {} err: {:?}", sql, e
        )))?;
        return match stmt {
            Statement::Statement(inner_stmt) => {
                self.dispatch_stmt(inner_stmt).await
            }
            Statement::CreateExternalTable(_) => {
                Err(Error::ExecutionErr("CreateExternalTable not support".into()))
            }
            Statement::CopyTo(_) => {
                Err(Error::ExecutionErr("ExecutionErr not support".into()))
            }
            Statement::Explain(_) => {
                Err(Error::ExecutionErr("Explain not support".into()))
            }
        };
    }

    async fn show_tables(&self, extended: bool, full: bool, db_name: &Option<Ident>, filter: &Option<ShowStatementFilter>) -> LResult<RecordBatch>;

    async fn drop(&self, object_type: &ObjectType, if_exists: bool, name: &Vec<ObjectName>) -> LResult<RecordBatch>;

    async fn create_database(&self, db_name: &ObjectName, if_not_exists: bool) -> LResult<RecordBatch>;

    fn session_context(&self) -> &SessionContext;

    fn post_query(&self, df: DataFrame) -> LResult<DataFrame>;

    async fn dispatch_stmt(&self, stmt: Box<SQLStatement>) -> LResult<Data> {
        use SQLStatement::*;
        match stmt.as_ref() {
            Drop { object_type, if_exists, names, .. } => {
                self.drop(object_type, *if_exists, names).await.map(Either::Right)
            }
            ShowTables { extended, full, db_name, filter } => {
                self.show_tables(*extended, *full, db_name, filter).await.map(Either::Right)
            }
            CreateDatabase { db_name, if_not_exists, .. } => {
                self.create_database(db_name, *if_not_exists).await.map(Either::Right)
            }
            Query(_) => {
                let state = self.session_context().state();
                let logic_plan = state
                    .statement_to_plan(Statement::Statement(stmt))
                    .await
                    .map_err(|e| Error::ExecutionErr(format!("create logic plan err: {:?}", e)))?;
                let best_plan = state.optimize(&logic_plan).map_err(|e| {
                    Error::ExecutionErr(format!("optimize err: {:?}", e))
                })?;
                let df = self.session_context().execute_logical_plan(best_plan).await.map_err(|e| {
                    Error::ExecutionErr(format!("execution err: {:?}", e))
                })?;
                self.post_query(df).map(Either::Left)
            }
            _ => {
                todo!()
            }
        }
    }
}

pub fn data_schema(data: &Data) -> SchemaRef {
    match data {
        Data::Left(df) => {
            Arc::new(df.schema().into())
        }
        Data::Right(r) => {
            r.schema()
        }
    }
}

